import cv2 as cv
import numpy as np
import torch
from torch import nn
from LeNet5 import LeNet5

def preprocess_image(image):
    """图像预处理"""
    gray = cv.cvtColor(image, cv.COLOR_BGR2GRAY)
    _, binary = cv.threshold(gray, 0, 255, cv.THRESH_BINARY_INV + cv.THRESH_OTSU)
    kernel = np.ones((3, 3), np.uint8)
    binary = cv.morphologyEx(binary, cv.MORPH_OPEN, kernel)
    return binary

def segment_digits(binary, original_image):
    """数字分割"""
    contours, _ = cv.findContours(binary, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)
    contours = sorted(contours, key=lambda x: cv.boundingRect(x)[0])

    digits = []
    for contour in contours:
        x, y, w, h = cv.boundingRect(contour)
        if w < 10 or h < 10:
            continue

        digit = binary[y:y + h, x:x + w]
        side = max(w, h)
        padding_top = (side - h) // 2
        padding_bottom = side - h - padding_top
        padding_left = (side - w) // 2
        padding_right = side - w - padding_left
        digit = cv.copyMakeBorder(digit, padding_top, padding_bottom, padding_left, padding_right,
                                  cv.BORDER_CONSTANT, value=0)

        digit = cv.resize(digit, (28, 28), interpolation=cv.INTER_AREA)

        if len(digit.shape) == 3:
            digit = cv.cvtColor(digit, cv.COLOR_BGR2GRAY)

        digits.append(digit)
        cv.rectangle(original_image, (x, y), (x + w, y + h), (0, 255, 0), 2)

    return digits

def recognize_digit(model, digit):
    """识别单个数字"""
    x = torch.from_numpy(digit).float() / 255.0
    x = x.reshape(1, 1, 28, 28)
    with torch.no_grad():
        logits = model(x)
        pred = logits.argmax(dim=1).item()
    return pred

def main(filepath):
    # 读取图片
    image = cv.imread(filepath)
    if image is None:
        print("无法读取图片")
        return

    # 保存原始图片的副本用于绘制边界框
    original_image = image.copy()

    # 预处理
    binary = preprocess_image(image)

    # 分割数字
    digits = segment_digits(binary, original_image)

    # 加载模型
    model = LeNet5()
    model.load_state_dict(torch.load('best.mdl', map_location=torch.device('cpu')))
    model.eval()

    # 识别数字
    student_id = ''
    for digit in digits:
        pred = recognize_digit(model, digit)
        student_id += str(pred)

    # 显示结果
    print('识别的学号为：', student_id)

    # 显示处理后的图片
    cv.imshow('Original Image with Boxes', original_image)
    cv.imshow('Processed Image', binary)
    cv.waitKey(0)
    cv.destroyAllWindows()

if __name__ == '__main__':
    main("./test/2.jpg")
